package com.lft.tree07.huffman_coding;

import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 哈夫曼编码常用操作
 */
public class HuffmanCoding {
	public static void main(String[] args) {
		// // 编码 —— 压缩 —— 分步调用
		// String content = "i like like like java do you like a java";
		// // 将字符串，转换为字节码数组。
		// byte[] contentBytes = content.getBytes();
		// // System.out.println(contentBytes.length); //40
		// // System.out.println(getNodes(contentBytes));
		//
		// // 将字节数组，转成字符对应数量的Node集合。
		// List<Node> nodes = getNodes(contentBytes);
		//
		// // 根据集合创建哈夫曼树。
		// Node node = createHuffmanTree(nodes);
		//
		// // 前序遍历
		// preOrder(node);
		// System.out.println("==============");
		// // 生成哈夫曼树对应的哈夫曼编码集合。
		// createHuffmanCodes(node);
		// System.out.println(map);
		// System.out.println("==============");
		// // 遍历 Map
		// for (Map.Entry<Byte, String> entry : map.entrySet()) {
		// 	System.out.println(
		// 			entry.getKey() + " => " + (char) (int) entry.getKey() + " => " + entry.getValue());
		// }
		// System.out.println("=============");
		// // 再根据哈夫曼编码集合，把原始的字节码数组转换为哈夫曼压缩后的字节码数组。
		// byte[] bytes = zipToHuffmanCodesBytes(contentBytes, map);
		// System.out.println(Arrays.toString(bytes));
		
		// // 原始字符串。
		// // String content = "i like like like java do you like a java";
		// String content = "你是谁？";
		// // 将字符串，转换为字节码数组。
		// byte[] contentBytes = content.getBytes();
		// // // 编码 —— 压缩
		// // 调用封装的方法生成经过哈夫曼编码压缩后的字节码数组
		// byte[] huffmanCodesBytes = getHuffmanCodesBytes(contentBytes);
		//
		// System.out.println(Arrays.toString(huffmanCodesBytes));
		//
		// // // 解码 —— 解压
		// // byteToBitString((byte) -1);
		// byte[] bytes = decode(huffmanCodes, huffmanCodesBytes);
		// System.out.println(new String(bytes));
		
		// // // 文件压缩
		zipFile("D://01.jpg", "D://03.zip");
		
		// // 文件解压
		unZipFile("D://03.zip", "D://04.jpg");
		
	}
	
	/**
	 * 哈夫曼编码步骤 2 ——【将字节码数组转成 Node 集合】
	 * 将字节数组，转成字符对应数量的Node集合。
	 * @param contentByteArray 字节数组
	 * @return 返回 List => [Node[data=97,weight=5],Node[data=32,weight=9]…]
	 */
	private static List<Node> byteArrayToNodeList(byte[] contentByteArray) {
		// 1. 创建一个 ArrayList
		List<Node> nodes = new ArrayList<>();
		
		// 2. 存储每个 byte出现的次数 -> map
		Map<Byte, Integer> map = new HashMap<>();
		for (byte b : contentByteArray) {
			Integer counter = map.get(b);
			// 未存放
			if (counter == null) {
				// 直接放进去。
				map.put(b, 1);
			} else {
				// 数量加1
				map.put(b, counter + 1);
			}
		}
		// 3. 将每一个键值对，转成一个Node对象，并加入到ArrayList
		for (Map.Entry<Byte, Integer> entry : map.entrySet()) {
			nodes.add(new Node(entry.getKey(), entry.getValue()));
		}
		return nodes;
	}
	
	/**
	 * 哈夫曼编码步骤 3 ——【将 Node 集合 创建哈夫曼树，返回树的根节点】
	 * 根据集合创建哈夫曼树。
	 * @param nodeList 原始集合。
	 * @return 返回哈夫曼树的根节点
	 */
	private static Node createHuffmanTree(List<Node> nodeList) {
		while (nodeList.size() > 1) {
			// 排序，从小到大
			Collections.sort(nodeList);
			// 取出第一棵最小的二叉树。
			Node leftNode = nodeList.get(0);
			// 取出第二棵次小的二叉树。
			Node rightNode = nodeList.get(1);
			// 生成一个新的二叉树，它的根节点没有 data，只有权限
			Node parent = new Node(null, leftNode.getWeight() + rightNode.getWeight());
			parent.setLeft(leftNode);
			parent.setRight(rightNode);
			
			// 将处理过的二叉树从 nodeList 中移除。
			nodeList.remove(leftNode);
			nodeList.remove(rightNode);
			
			// 将新的二叉树添加进集合。
			nodeList.add(parent);
		}
		// nodeList 中最后的结点，就是哈夫曼树的根节点。
		return nodeList.get(0);
	}
	
	// 1. 用于存放哈夫曼编码
	static Map<Byte, String> huffmanCodes = new HashMap<>();
	// // 2. 在生成哈夫曼编码表时，需要去拼接路径，定义一个 StringBuilder 存储某个叶子节点的路径。
	static StringBuilder globalStringBuilder = new StringBuilder();
	
	/**
	 * 哈夫曼编码步骤 4 ——【根据 哈夫曼树的根节点，生成哈夫曼编码集合[字节码，"1001"]】
	 * 根据传入的哈夫曼树的根节点，生成哈夫曼树对应的哈夫曼编码集合。
	 * 1. 将哈夫曼编码表，存放在 Map<Byte,String>集合中。
	 * //  形式：[(32,01),(97,100)...]
	 * 2. 在生成哈夫曼编码表时，需要去拼接路径，定义一个 StringBuilder 存储某个叶子节点的路径。
	 * StringBuilder stringBuilder = new StringBuilder();
	 * @param root 传入的哈夫曼根节点。
	 * @param code 路径：左子节点0，右子节点1
	 * @return 返回 集合。
	 */
	private static Map<Byte, String> createHuffmanCodes(Node root, String code, StringBuilder stringBuilder) {
		if (root == null) {
			throw new RuntimeException("传入的根节点为空！");
		}
		// 2. 在生成哈夫曼编码表时，需要去拼接路径，定义一个 StringBuilder 存储某个叶子节点的路径。
		StringBuilder newStringBuilder = new StringBuilder(stringBuilder);
		// 用于拼接路径
		newStringBuilder.append(code);
		// 没有数据的节点，是非叶子节点
		if (root.getData() == null) {
			// 递归处理
			// 向左递归
			createHuffmanCodes(root.getLeft(), "0", newStringBuilder);
			// 向右递归
			createHuffmanCodes(root.getRight(), "1", newStringBuilder);
		} else {
			// 找到叶子节点
			huffmanCodes.put(root.getData(), newStringBuilder.toString());
		}
		return huffmanCodes;
	}
	
	/**
	 * 哈夫曼编码步骤 5 ——【根据 哈夫曼树的根节点，生成哈夫曼编码集合[字节码，"1001"]】
	 * 方便调用，重载 createHuffmanCodes 方法。
	 * @param root 传入哈夫曼树根节点
	 * @return
	 */
	private static Map<Byte, String> createHuffmanCodes(Node root) {
		if (root == null) {
			return null;
		}
		// 处理左子树
		createHuffmanCodes(root.getLeft(), "0", globalStringBuilder);
		// 处理左子树
		createHuffmanCodes(root.getRight(), "1", globalStringBuilder);
		return huffmanCodes;
	}
	
	/**
	 * 哈夫曼编码步骤 6 ——【根据哈夫曼编码集合，将原始字节码数组 转换压缩成 哈夫曼编码压缩后的字节码数组】
	 * 将字符串对应的 byte[] 数组 通过生成的哈夫曼编码表，返回一个哈夫曼编码压缩后的 byte[] 数组
	 * 举例：String content = "i like 123 321 like like java do you like a java";
	 * => byte[] contentBytes = content.getBytes();
	 * =>1010100010111111110010001011111111001000101111111100100101001101110001110000011011101000111100101000101111111100110001001010011011100
	 * =>对应的 byte[] zipToHuffmanCodesBytes，即 8 位对应一个byte,放入 zipToHuffmanCodesBytes 数组中
	 * => zipToHuffmanCodesBytes[0] = 10101000(补码) => 补码 -1 => 10100111(反码) => 符号位不变，其他位取反 => 11011000(原码) => -88
	 * @param contentBytes 原始字符串对应的 byte[] 数组
	 * @param huffmanCodes 提供哈夫曼编码表。
	 * @return 返回哈夫曼编码处理后的 byte[] 数组
	 */
	private static byte[] toHuffmanCodesByteArrays(byte[] contentBytes, Map<Byte, String> huffmanCodes) {
		// 1. 利用 huffmanCodes 将 contentBytes 转成 哈夫曼编码对应的字符串。
		StringBuilder stringBuilder = new StringBuilder();
		for (byte b : contentBytes) {
			stringBuilder.append(huffmanCodes.get(b));
		}
		// 将 stringBuilder 字符串，转换成 byte[] 数组
		int length;
		if (stringBuilder.length() % 8 == 0) {
			length = stringBuilder.length() / 8;
		} else {
			length = stringBuilder.length() / 8 + 1;
		}
		
		// 创建存储压缩后的 byte 数组
		byte[] huffmanCodeBytes = new byte[length];
		// 记录是第几个 byte
		int index = 0;
		// 因为是每 8 位对应一个 byte，所以步长是 8
		for (int i = 0; i < stringBuilder.length(); i += 8) {
			String strByte;
			if (i + 8 > stringBuilder.length()) {
				strByte = stringBuilder.substring(i);
			} else {
				strByte = stringBuilder.substring(i, i + 8);
			}
			// 将 strByte 以二进制形式转成 byte[]
			huffmanCodeBytes[index++] = (byte) Integer.parseInt(strByte, 2);
		}
		return huffmanCodeBytes;
	}
	
	/**
	 * 哈夫曼编码步骤 7 ——【封装压缩方法】
	 * 【哈夫曼编码——压缩】
	 * 封装，由原始的字节码数组转换成哈夫曼编码压缩后的字节码数组
	 * @param content 原始的字节码数组。
	 * @return 哈夫曼编码压缩全的字节码数组
	 */
	public static byte[] getHuffmanCodesBytes(byte[] contentByteArray) {
		if (contentByteArray == null) {
			return null;
		}
		// 将字节数组，转成字符对应数量的Node集合。
		List<Node> nodes = byteArrayToNodeList(contentByteArray);
		// 根据节点的集合创建哈夫曼树。
		Node node = createHuffmanTree(nodes);
		// 生成哈夫曼树对应的哈夫曼编码集合。
		Map<Byte, String> huffmanMap = createHuffmanCodes(node);
		// 再根据哈夫曼编码集合，把原始的字节码数组转换为哈夫曼压缩后的字节码数组，再返回。
		return toHuffmanCodesByteArrays(contentByteArray, huffmanMap);
	}
	
	/**
	 * 哈夫曼编码步骤 8 ——【解压步骤1——将一个字节码转换成一个8位的字符串,最后一段不满8位的不需要补高位】
	 * 1. 将压缩后的字节码数组 [-88, -65, -56, -65, -56, -65, -55, 77, -57, 6, -24, -14, -117, -4, -60, -90, 28]
	 * 转成 哈夫曼编码对应的进制的字符串
	 * 1010100010111111110010001011111111001000101111111100100101001101110001110000011011101000111100101000101111111100110001001010011011100
	 * 2. 对照 哈夫曼编码 再转成 字符串 “i like like like java do you like a java”
	 * @param b    需要转成二进制的1个字节码
	 * @param flag 标志是否需要补高位，true 需要补高位，false 表示不需要补高位。如果是最后一个字节，无需补高位。
	 * @return 是该 b 对应的二进制的字符串。（注意是按补码返回）
	 */
	private static String byteToBitString(byte b, boolean flag) {
		// 使用变量保存b，将 b 转换成 int(32位)
		int temp = b;
		if (flag) {
			// 如果是正数我们还存在补高位 => 按位或
			// 256 | 1 -> 1 0000 0000 | 0000 0001 -> 1 0000 0001
			temp |= 256;
		}
		// 将 int 转换成 二进制
		String str = Integer.toBinaryString(temp);
		
		if (flag) {
			// 截取最后 8 位
			return str.substring(str.length() - 8);
		} else {
			return str;
		}
	}
	
	/**
	 * 哈夫曼编码步骤 9 ——【解压步骤2——将压缩后的哈夫曼编码字节码数组 能过 哈夫曼编码表 解压成 原始内容字节码数组】
	 * @param huffmanCodes 哈夫曼编码表
	 * @param huffmanBytes 哈夫曼字编码压缩过的节数组
	 * @return
	 */
	private static byte[] decode(Map<Byte, String> huffmanCodes, byte[] huffmanBytes) {
		// 1. 先得到 huffmanBytes 字节码数组 对应的 二进制字符串。形式：10101000101111111100100010...
		StringBuilder stringBuilder = new StringBuilder();
		
		// 2. 将 byte[] 数组转成二进制的字符串
		for (int i = 0; i < huffmanBytes.length; i++) {
			stringBuilder.append(byteToBitString(huffmanBytes[i], i == huffmanBytes.length - 1 ? false : true));
		}
		
		// System.out.println("二进制字符串: " + stringBuilder);
		
		// 把字符串按照指定的哈夫曼编码表进行解码
		// 将哈夫曼编码表进行调换。反向查询 a -> 97 => 97 -> a
		Map<String, Byte> decodeHuffmanCodes = new HashMap<>();
		// 遍历 原来的哈夫曼编码表。
		for (Map.Entry<Byte, String> entry : huffmanCodes.entrySet()) {
			decodeHuffmanCodes.put(entry.getValue(), entry.getKey());
		}
		
		// 反转后的哈夫曼解码表。
		// System.out.println("哈夫曼解码表: " + decodeHuffmanCodes);
		
		// 创建一个集合，存放 byte
		List<Byte> list = new ArrayList<>();
		// 遍历 二进制 对应的字符串
		// 用于指定子串的起始位置。每次从该位置找，直到找到一个能在map中找到的。
		int temp = 0;
		// 【注意】i从1开始，substring是从 temp 开始，取 i-temp 个
		for (int i = 1; i <= stringBuilder.length(); i++) {
			String substring = stringBuilder.substring(temp, i);
			Byte aByte = decodeHuffmanCodes.get(substring);
			if (aByte != null) {
				list.add(aByte);
				temp = i;
			}
		}
		
		// 循环结束后。list 中保存了所有的字符。
		// System.out.println("解码后的每个字节码 " + list);
		
		byte[] bytes = new byte[list.size()];
		for (int i = 0; i < list.size(); i++) {
			bytes[i] = list.get(i);
		}
		// System.out.println("解码之后得到的byte[]数组： " + Arrays.toString(bytes));
		return bytes;
	}
	
	/**
	 * 哈夫曼编码步骤 10 ——【压缩文件】
	 * @param srcFileAddress 要压缩源文件路径
	 * @param dstFileAddress 压缩文件保存路径
	 */
	private static void zipFile(String srcFileAddress, String dstFileAddress) {
		// 创建文件的输入流
		FileInputStream fis = null;
		OutputStream os = null;
		ObjectOutputStream oos = null;
		try {
			fis = new FileInputStream(srcFileAddress);
			// 创建一个和源文件大小一样的byte[]
			byte[] bytes = new byte[fis.available()];
			// 将流中的字节码，读取到 创建的字节码数组中。
			fis.read(bytes);
			// 直接对源文件进行压缩。
			byte[] huffmanCodesBytes = getHuffmanCodesBytes(bytes);
			
			// 创建文件的输出流。
			os = new FileOutputStream(dstFileAddress);
			
			// 创建一个和文件输出流关联的ObjectOutputStream
			oos = new ObjectOutputStream(os);
			
			// 先把哈夫曼编码后的字节数组写入压缩文件。
			oos.writeObject(huffmanCodesBytes);
			// 这里我们以对象流的方式写入哈夫曼编码。是为了以后我们恢复源文件时使用。
			// 【注意】一定要把哈夫曼编码写入压缩文件。不然回头解压的时候没有编码
			oos.writeObject(huffmanCodes);
			System.out.println("文件压缩成功！");
		} catch (Exception e) {
			e.printStackTrace();
		} finally {
			if (oos != null) {
				try {
					oos.close();
				} catch (IOException e) {
					e.printStackTrace();
				}
			}
			if (os != null) {
				try {
					os.close();
				} catch (IOException e) {
					e.printStackTrace();
				}
			}
			if (fis != null) {
				try {
					fis.close();
				} catch (IOException e) {
					e.printStackTrace();
				}
			}
		}
	}
	
	/**
	 * 哈夫曼编码步骤 10 ——【解压文件】
	 * @param srcFileAddress 压缩文件路径
	 * @param dstFileAddress 解压后文件路径
	 */
	private static void unZipFile(String srcFileAddress, String dstFileAddress) {
		// 定义文件输入流
		InputStream is = null;
		// 定义一个对象输入流
		ObjectInputStream ois = null;
		// 定义文件输出流
		OutputStream os = null;
		try {
			// 创建文件输入流
			is = new FileInputStream(srcFileAddress);
			// 创建一个和 is 关联的对象输入流。
			ois = new ObjectInputStream(is);
			
			// 读取 byte[] 数组，huffmanCodesBytes
			byte[] huffmanCodesBytes = (byte[]) ois.readObject();
			
			// 读取 哈夫曼编码表，huffmanCodes
			Map<Byte, String> huffmanCodes = (Map<Byte, String>) ois.readObject();
			
			// 解压
			byte[] bytes = decode(huffmanCodes, huffmanCodesBytes);
			
			// 创建输出流
			os = new FileOutputStream(dstFileAddress);
			// 写入到目标文件中。
			os.write(bytes);
			System.out.println("文件解压成功！");
		} catch (Exception e) {
			e.printStackTrace();
		} finally {
			if (os != null) {
				try {
					os.close();
				} catch (IOException e) {
					e.printStackTrace();
				}
			}
			if (ois != null) {
				try {
					ois.close();
				} catch (IOException e) {
					e.printStackTrace();
				}
			}
			if (is != null) {
				try {
					is.close();
				} catch (IOException e) {
					e.printStackTrace();
				}
			}
		}
	}
	
	/**
	 * 前序遍历哈夫曼树
	 * @param node 传入哈夫曼树的根节点
	 */
	private static void preOrder(Node node) {
		if (node == null) {
			System.out.println("树为空，无法遍历！");
		} else {
			node.preOrder();
		}
	}
	
}

/**
 * 哈夫曼编码步骤 1 ——【创建节点类】
 * 创建 Node，带数据和权值
 */
class Node implements Comparable<Node> {
	// 存放数据本身，比如 'a' => 97; ' ' => 32
	private Byte data;
	// 权值，表示字符出现的次数
	private int weight;
	// 指向节点的左子节点
	private Node left;
	// 指向节点的右子节点
	private Node right;
	
	public Node(Byte data, int weight) {
		this.data = data;
		this.weight = weight;
	}
	
	public Byte getData() {
		return data;
	}
	
	public void setData(Byte data) {
		this.data = data;
	}
	
	public int getWeight() {
		return weight;
	}
	
	public void setWeight(int weight) {
		this.weight = weight;
	}
	
	public Node getLeft() {
		return left;
	}
	
	public void setLeft(Node left) {
		this.left = left;
	}
	
	public Node getRight() {
		return right;
	}
	
	public void setRight(Node right) {
		this.right = right;
	}
	
	/**
	 * 前序遍历
	 */
	public void preOrder() {
		System.out.println(this);
		if (this.getLeft() != null) {
			this.getLeft().preOrder();
		}
		if (this.getRight() != null) {
			this.getRight().preOrder();
		}
	}
	
	@Override
	public String toString() {
		return "Node [" +
				"Data=" + data +
				", Weight=" + weight +
				']';
	}
	
	@Override
	public int compareTo(Node o) {
		return this.getWeight() - o.getWeight();
	}
}